import socket
import os
import struct
import datetime

PORT = 4321
MAX_FILENAME = 256
BUF_SIZE = 4096

def recv_all(conn, length):
    data = b''
    while len(data) < length:
        try:
            chunk = conn.recv(length - len(data))
        except Exception as e:
            print(f"接收异常: {e}")
            return None
        if chunk is None:
            return None
        if chunk == b'':
            # 只有真正断开才返回None
            return None
        data += chunk
    return data

def main():
    if not os.path.exists('photos'):
        os.makedirs('photos')

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(('0.0.0.0', PORT))
    sock.listen(1)
    print(f"等待连接在端口{PORT}...")

    try:
        while True:
            conn, addr = sock.accept()
            print(f"连接来自: {addr[0]}")
            try:
                while True:
                    # 1. 接收4字节（文件名长度或心跳包）
                    header = recv_all(conn, 4)
                    if header is None:
                        print("客户端断开连接，关闭连接")
                        break
                    if len(header) != 4:
                        print("收到无效头，忽略")
                        continue
                    # 检查心跳包
                    if header == b'PING':
                        try:
                            conn.sendall(b'PONG')
                        except Exception as e:
                            print(f"回复PONG失败: {e}")
                            break
                        print("收到心跳包，已回复PONG")
                        continue
                    # 文件名长度
                    filename_size = struct.unpack('<I', header)[0]
                    print(f"接收文件名大小: {filename_size}字节")
                    if filename_size <= 0 or filename_size > 255:
                        print(f"Invalid filename size: {filename_size}")
                        continue
                    # 文件名
                    filename_bytes = recv_all(conn, filename_size)
                    if filename_bytes is None or len(filename_bytes) != filename_size:
                        print("文件名接收失败，忽略")
                        continue
                    filename = filename_bytes.decode('utf-8', errors='ignore').strip()
                    if len(filename) < 4 or not filename.endswith('.jpg'):
                        print("Invalid file extension")
                        continue
                    # 生成以时间为名的文件名
                    now = datetime.datetime.now()
                    time_filename = now.strftime("photos/%Y%m%d_%H%M%S_%f")[:-3] + ".jpg"
                    print(f"保存文件名: {time_filename}")
                    # 文件大小
                    size_buf = recv_all(conn, 4)
                    if size_buf is None or len(size_buf) != 4:
                        print("Invalid file size header，忽略")
                        continue
                    file_size = struct.unpack('<I', size_buf)[0]
                    print(f"接收文件大小: {file_size}字节")
                    if file_size <= 0:
                        print(f"Invalid file size: {file_size}")
                        continue
                    # 文件内容
                    received = 0
                    os.makedirs('photos', exist_ok=True)  # 确保目录存在
                    with open(time_filename, 'wb') as f:
                        while received < file_size:
                            to_read = min(BUF_SIZE, file_size - received)
                            chunk = conn.recv(to_read)
                            if not chunk:
                                print("文件内容接收中断")
                                break
                            f.write(chunk)
                            received += len(chunk)
                            print(f"已接收: {received}/{file_size}字节", end='\r')
                        f.flush()
                        os.fsync(f.fileno())
                    print(f"\n文件保存成功: {time_filename}")
            finally:
                conn.close()
                print("连接关闭")
    except KeyboardInterrupt:
        print("\n服务器关闭")
    finally:
        sock.close()

if __name__ == '__main__':
    main()